#ifndef HELLOWORLD_DXGI_H
#define HELLOWORLD_DXGI_H
#include <d3d11.h>
#include <dxgi1_2.h>
#include <stdio.h>
#include <ctime>
#include <iostream>
#include <Windows.h>


#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui_c.h>
#include <opencv2/highgui/highgui.hpp>

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "dxgi.lib")

#pragma warning(disable:4996)



#endif //HELLOWORLD_DXGI_H


class DXGICapture
{
public:
    DXGICapture();
    ~DXGICapture();

public:
    bool Update();
    bool BitmapToShow();
    void SetFps(int fps);

private:
    HRESULT hr; // 初始化函数返回状态
    IDXGIOutputDuplication *_pDXGIOutputDup = nullptr;
    IDXGIResource *desktopResource = nullptr;
    ID3D11Texture2D *_pDX11Texture = nullptr;
    ID3D11Texture2D *_pCopyBuffer = nullptr;
    IDXGISurface *CopySurface = nullptr;
    ID3D11Device *_pDX11Dev = nullptr;
    ID3D11DeviceContext *_pDX11DevCtx = nullptr;
    DXGI_MAPPED_RECT MappedSurface;

    int fps;

private:
    bool Init();
    bool RGBDataSaveAsBmp(
            unsigned char * pRgbData,            // 图像数据
            int width,                           // 图像宽度
            int height,                          // 图像高度
            int biBitCount,                      // 位图深度
            bool flipvertical);

};

DXGICapture::DXGICapture() {
    Init();
};

DXGICapture::~DXGICapture() {
    //释放资源
};

bool DXGICapture::Init() {
    // Driver types supported 支持的驱动程序类型
    D3D_DRIVER_TYPE DriverTypes[] = {
            D3D_DRIVER_TYPE_HARDWARE,
            D3D_DRIVER_TYPE_WARP,
            D3D_DRIVER_TYPE_REFERENCE,
    };
    // 计算有几类数据驱动类型
    UINT NumDriverTypes = ARRAYSIZE(DriverTypes);

    // 支持的功能级别
    D3D_FEATURE_LEVEL FeatureLevels[] =
            {
                    D3D_FEATURE_LEVEL_11_0,
                    D3D_FEATURE_LEVEL_10_1,
                    D3D_FEATURE_LEVEL_10_0,
                    D3D_FEATURE_LEVEL_9_1
            };
    // 支持功能级别的数量
    UINT NumFeatureLevels = ARRAYSIZE(FeatureLevels);

    D3D_FEATURE_LEVEL FeatureLevel;

    IDXGIOutput1 *_pDXGIOutput1 = nullptr;

    // 找到匹配的驱动后调出循环
    for(UINT index = 0; index < NumDriverTypes; index++)
    {
        hr = D3D11CreateDevice(
                nullptr,
                DriverTypes[index],
                nullptr,0,
                FeatureLevels,
                NumFeatureLevels,
                D3D11_SDK_VERSION,
                &_pDX11Dev,
                &FeatureLevel,
                &_pDX11DevCtx);

        if (SUCCEEDED(hr)){
            break;
        }
    }
    IDXGIDevice *_pDXGIDev = nullptr;
    // Get DXGI device 获取 DXGI 设备
    hr = _pDX11Dev->QueryInterface(__uuidof(IDXGIDevice), reinterpret_cast<void**>(&_pDXGIDev));
    if (FAILED(hr)) {
        return false;
    }

    IDXGIAdapter *_pDXGIAdapter = nullptr;
    // Get DXGI adapter 获取 DXGI 适配器
    hr = _pDXGIDev->GetParent(__uuidof(IDXGIAdapter), reinterpret_cast<void**>(&_pDXGIAdapter));
    if (FAILED(hr)) {
        return false;
    }

    UINT i = 0;
    IDXGIOutput *_pDXGIOutput = nullptr;
    // Get output 获取输出
    hr = _pDXGIAdapter->EnumOutputs(i, &_pDXGIOutput);
    if (FAILED(hr)) {
        return false;
    }

    DXGI_OUTPUT_DESC DesktopDesc;
    // Get output description struct 获取输出描述结构
    _pDXGIOutput->GetDesc(&DesktopDesc);

    // 判断是否可以正常从接口获取数据，如果失败调出。
    hr = _pDXGIOutput->QueryInterface(__uuidof(IDXGIOutput1), reinterpret_cast<void**>(&_pDXGIOutput1));
    if (FAILED(hr)) {
        return false;
    }

    // Create desktop duplication 创建桌面副本
    hr = _pDXGIOutput1->DuplicateOutput(_pDX11Dev, &_pDXGIOutputDup);
    if (FAILED(hr)) {
        return false;
    }

    return true;
}

bool DXGICapture::Update(){
    DXGI_OUTDUPL_FRAME_INFO frameInfo;
    hr = _pDXGIOutputDup->AcquireNextFrame(1000, &frameInfo, &desktopResource);
    if (FAILED(hr))
    {
        if (hr == DXGI_ERROR_WAIT_TIMEOUT)
        {
            if (desktopResource) {
                desktopResource->Release();
                desktopResource = nullptr;
            }
            hr = _pDXGIOutputDup->ReleaseFrame();
        }
        else
        {
            return false;
        }
    }

    // query next frame staging buffer 查询下一帧暂存缓冲区
    hr = desktopResource->QueryInterface(__uuidof(ID3D11Texture2D), reinterpret_cast<void **>(&_pDX11Texture));
    desktopResource->Release();

    desktopResource = nullptr;
    if (FAILED(hr)) {
        return false;
    }

    D3D11_TEXTURE2D_DESC desc;
    // copy old description 复制旧描述
    if (_pDX11Texture)
    {
        _pDX11Texture->GetDesc(&desc);
    }
    else if (_pCopyBuffer)
    {
        _pCopyBuffer->GetDesc(&desc);
    }
    else
    {
        return false;
    }

    // create a new staging buffer for fill frame image 为填充帧图像创建一个新的暂存缓冲区
    if (_pCopyBuffer == nullptr) {
        D3D11_TEXTURE2D_DESC CopyBufferDesc;
        CopyBufferDesc.Width = desc.Width;
        CopyBufferDesc.Height = desc.Height;
        CopyBufferDesc.MipLevels = 1;
        CopyBufferDesc.ArraySize = 1;
        CopyBufferDesc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
        CopyBufferDesc.SampleDesc.Count = 1;
        CopyBufferDesc.SampleDesc.Quality = 0;
        CopyBufferDesc.Usage = D3D11_USAGE_STAGING;
        CopyBufferDesc.BindFlags = 0;
        CopyBufferDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        CopyBufferDesc.MiscFlags = 0;

        hr = _pDX11Dev->CreateTexture2D(&CopyBufferDesc, nullptr, &_pCopyBuffer);
        if (FAILED(hr)) {
            return false;
        }
    }

    if (_pDX11Texture)
    {
        // copy next staging buffer to new staging buffer 将下一个暂存缓冲区复制到新的暂存缓冲区
        _pDX11DevCtx->CopyResource(_pCopyBuffer, _pDX11Texture);
    }

    // create staging buffer for map bits 为映射位创建暂存缓冲区
    hr = _pCopyBuffer->QueryInterface(__uuidof(IDXGISurface), (void **)&CopySurface);
    if (FAILED(hr)) {
        return false;
    }


    // copy bits to user space 将位复制到用户空间
    hr = CopySurface->Map(&MappedSurface, DXGI_MAP_READ);


    return true;
};

bool DXGICapture::RGBDataSaveAsBmp(
        unsigned char *pRgbData,
        int width,
        int height,
        int biBitCount,
        bool flipvertical)
{

    cv::Mat dst_mat(cvSize(1920, 1080), CV_MAKETYPE(CV_8UC3, 4));
    memcpy(dst_mat.data, pRgbData, 1920*1080*4*sizeof(unsigned char));

    int font_face = cv::FONT_HERSHEY_COMPLEX;
    double font_scale = 1;
    int thickness = 2;
    cv::putText(dst_mat, std::to_string(fps)+" fps", cv::Point(80, 40), font_face, font_scale, cv::Scalar(0, 255, 255), thickness, 8, 0);
    cv::imshow("ScreenCapture",dst_mat);
    if(27 == cv::waitKey(1))
        return false;
    return true;
}

bool DXGICapture::BitmapToShow(){
    SetConsoleOutputCP(CP_UTF8);
    if(!RGBDataSaveAsBmp(this->MappedSurface.pBits,1920, 1080, 32, true))
        return false;
    CopySurface->Unmap();
    hr = CopySurface->Release();
    CopySurface = nullptr;

    if (_pDXGIOutputDup)
    {
        hr = _pDXGIOutputDup->ReleaseFrame();
    }
    return true;
}

void DXGICapture::SetFps(int fps){
    this->fps = fps;
}
